package com.eonmind.ai.mapper;

import com.eonmind.ai.dto.AiDrawDto;
import com.eonmind.common.contract.data.base.BaseMongoMapper;
import com.eonmind.common.contract.data.request.PageRequest;
import com.eonmind.common.contract.enums.DelEnum;
import com.eonmind.common.contract.enums.ai.AiArtworkStatusEnum;
import com.eonmind.common.provider.entity.ai.AiArtwork;
import com.eonmind.user.contract.dto.UserInfoData;
import com.mongoplus.conditions.query.QueryWrapper;
import com.mongoplus.model.PageResult;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class AiDrawMapper extends BaseMongoMapper<AiArtwork> {


    public AiArtwork selectById(Long id) {
        QueryWrapper<AiArtwork> eq = createMgoQuery()
                .eq(AiArtwork::getId, id)
                .eq(AiArtwork::getDel, DelEnum.UN_DEL.getCode());
        return one(eq);
    }

    public List<AiArtwork> selectByStatus(Integer code) {
        QueryWrapper<AiArtwork> eq = createMgoQuery()
                .eq(AiArtwork::getStatus, code)
                .eq(AiArtwork::getDel, DelEnum.UN_DEL.getCode());
        return list(eq);
    }

    public List<AiArtwork> selectAllSyncData() {
        QueryWrapper<AiArtwork> mgoQuery = createMgoQuery().or(t -> {
            return t.eq(AiArtwork::getStatus, AiArtworkStatusEnum.PENDING_GENERATION.getCode())
                    .eq(AiArtwork::getStatus, AiArtworkStatusEnum.GENERATING.getCode());
        });

        mgoQuery.eq(AiArtwork::getDel, DelEnum.UN_DEL.getCode());
        return list(mgoQuery);
    }

    public PageResult<AiArtwork> pageDraw(UserInfoData userInfoData, AiDrawDto.AiDrawPageRequest aiDrawPageRequest, Integer platformType) {
        QueryWrapper<AiArtwork> eq = createMgoQuery()
                .eq(AiArtwork::getUserId, userInfoData.getId())
                .eq(AiArtwork::getPlatformType, platformType)
                .eq(AiArtwork::getDel, DelEnum.UN_DEL.getCode());
        PageRequest page = aiDrawPageRequest.getPage();
        return page(eq, page);
    }

    public List<AiArtwork> selectUserWaitTask(Long userId) {
        QueryWrapper<AiArtwork> or = createMgoQuery()
                .eq(AiArtwork::getUserId, userId)
                .or(t -> {
                    return t.eq(AiArtwork::getStatus, AiArtworkStatusEnum.PENDING_GENERATION.getCode())
                            .eq(AiArtwork::getStatus, AiArtworkStatusEnum.GENERATING.getCode());
                });
        return list(or);
    }
}
